from typing import Optional, Any, Iterator
from openai import OpenAI
from zai import ZhipuAiClient
from langchain.llms.base import LLM
from langchain_core.outputs import GenerationChunk
import os


# base_url = 'http://localhost:8000/v1/'


class GLM_LLMS(LLM):
    client: ZhipuAiClient = ZhipuAiClient(api_key="d0fc2026b50344b18e25187d9393ce3f.P2XsXy1lpeqc2Gl0")
    # client: OpenAI = OpenAI(api_key='EMPTY', base_url=base_url)
    loacl_model: str = "GLM-4-Flash-250414"

    # loacl_model: str = "LOCAL_GLM-3-6B"

    def __init__(self):
        super().__init__()

    @property
    def _llm_type(self) -> str:
        return "GLM-4-Flash-250414"

    def _call(
            self,
            prompt: str,
            stop: Optional[list[str]] = None,
            run_manager=None,
            **kwargs: Any,
    ) -> str:
        messages = [{"role": "user", "content": prompt}]
        response = self.client.chat.completions.create(
            model=self.loacl_model,
            messages=messages,
        )
        # print("_call ====response=====> ", response)
        return response.choices[0].message.content

    def _stream(
            self,
            prompt: str,
            stop: Optional[list[str]] = None,
            run_manager=None,
            **kwargs: Any,
    ) -> Iterator:
        # print("_stream ===agent_output======> ", agent_output)
        messages = [
            # {"role": "system", "content": f"你可以参考代理返回的数据进行最后的回答。 代理返回如下:{agent_output}"},
            {"role": "user", "content": prompt}
        ]
        # 获取最终回答
        response = self.client.chat.completions.create(
            model=self.loacl_model,
            messages=messages,
            stream=True
        )
        append_text = ""
        for chunk in response:
            if chunk.choices[0].delta.content:
                append_text += chunk.choices[0].delta.content
                yield GenerationChunk(text=append_text)


from langchain.prompts import ChatPromptTemplate


def get_mind_chat(message, current_his):
    print("======get_mind_chat=======")
    llm = GLM_LLMS()
    template = """
    用户的问题：{message}
    如果用需要向生成一个思维导图， 你需要生成与问题相关的一个支持markmap的markdown文档，并且只返回这个markdown内容,去掉开头的```markdown  和结尾的```标记,如果和生成思维导图无关，你可以正常根据你的知识回答用户。
    """
    prompt = ChatPromptTemplate.from_template(template)
    chain = prompt | llm
    return chain.invoke({"message": message})


def get_mind_chat_steam(message, current_his):
    print("======get_mind_chat_steam=======")
    llm = GLM_LLMS()
    template = """
    用户的问题：{message}
    如果用需要向生成一个思维导图， 你需要生成与问题相关的一个支持markmap的markdown文档，并且只返回这个markdown内容，去掉开头的```markdown  和结尾的```标记,如果和生成思维导图无关，你可以正常根据你的知识回答用户。
    """
    prompt = ChatPromptTemplate.from_template(template)
    chain = prompt | llm
    # chain.stream 返回的是一个生成器，产生 GenerationChunk 对象
    # 我们需要将其转换为字符串生成器
    for chunk in chain.stream({"message": message}):
        yield chunk


def get_dir_list(path):
    res = [path + '\\' + file for file in os.listdir(path)]
    return res


def get_dir_list_str(path):
    res = [path + '\\' + file for file in os.listdir(path)]
    return res.__str__()

import torch

if __name__ == '__main__':
    # print(get_mind_chat("给我生成一个学习ai的思维导图"))
    # for item in get_mind_chat_steam("给我生成一个学习ai的思维导图"):
    #     print(item, end="", flush=True)
    print(torch.__version__)
    print(torch.cuda.is_available())